/*
 * Copyright (c) China Telecom Cloud Technology Co., Ltd. 2024-2025. All rights reserved.
 */
#include <iostream>
#include <sys/mman.h>
#include <sys/stat.h>
#include <ctime>
#include <fcntl.h>
#include <unistd.h>
#include <cstdint>
#include <cassert>
#include <sys/mman.h>
#include <unistd.h>
#include <cerrno>
#include <cstring>
#include <dirent.h>
#include <cstring>
#include "Common.h"
#include "MemoryPool.h"
#include "cstdio"
#include "Logger.h"


// TODO(yangxianpku: 2025.08.18): 运行时动态配置
//! 共享文件路径
std::string generateSharedDir() {
    return "/mnt/hugepages/";
}


//! 共享文件名前缀
std::string generateShareFilePrefix() {
    return "memorypool-";
}


//! KV Store服务的端口号
std::string generateClientPort() {
    return ConfigLoader::get_config<std::string>(
                ConfigLoader::EnvVars::CLIENT_PORT, 
                DEFAULT_CLIENT_PORT
            );
}


//! 生成共享文件名
// TODO(yangxianpku: 2025.08.18): 添加".kvstore"后缀
// (base) [yangxianpku@localhost hugepages]$ pwd
// /mnt/hugepages
// (base) [yangxianpku@localhost hugepages]$ ls -lah
// total 8.1G
// drwxr-xr-x. 2 root root    0 Sep  2 11:38 .
// drwxr-xr-x. 5 root root 4.0K Aug 29 09:55 ..
// -rw-r--r--. 1 root root 8.0G Sep  2 11:38 memorypool-50051-33554432-1756784289464490
std::string generateSharedFilename(uint32_t memSegSize) {
    return generateSharedDir()           + 
        generateShareFilePrefix()        + 
        generateClientPort()             + 
        "-" + std::to_string(memSegSize) + 
        "-" + std::to_string(getCurrentTimestampUs());
}



int32_t MemoryPool::getIdxByOffset(uint64_t offset) {
    return offset / _phyMemSegSize;
}

void MemoryPool::init(void *memStartAddr, 
                    uint32_t memSegSize, 
                    uint32_t phyMemSegSize,
                    uint64_t poolSize, 
                    uint64_t phyPoolSize, 
                    int32_t blkNum, 
                    bool isShare,
                    std::list<uint64_t> &freeList, 
                    std::string shmFilePath, 
                    std::vector<Segment> &segList) {
    auto start = getCurrentTimestampUs();  // 记录函数开始时间

    _memStartAddr  = memStartAddr;
    _memSegSize    = memSegSize;
    _phyMemSegSize = phyMemSegSize;
    _poolSize      = poolSize;
    _phyPoolSize   = phyPoolSize;
    _blkNum        = blkNum;
    _isShare       = isShare;
    _freeList      = freeList;
    _segList       = segList;
    _shmFilePath   = shmFilePath;

    // 重置KV Cache请求与命中统计
    resetCacheStatis();

    LOG_INFO_FMT("blkNum:{}, freeList.size:{}, segList.size:{}", 
                _blkNum, _freeList.size(), _segList.size());

    auto end = getCurrentTimestampUs();  // 记录函数结束时间
    LOG_DEBUG_FMT("cost time : {} us", end - start);
}


std::vector<Segment>& MemoryPool::getSeglist() {
    return _segList;
}


std::list<uint64_t>& MemoryPool::getFreelist() {
    return _freeList;
}


int32_t MemoryPool::allocateSegmentByList(
            std::vector<BlockHashId> &blockHashIds,  // input
            AllocPolicy policy,                
            std::vector<Segment *> &segments,   // out; [nullptr, .....]
            int32_t proposeNum, 
            int32_t writeInRecordSize, 
            CacheInfo cacheInfo) {
    auto start = getCurrentTimestampUs();  // 记录函数开始时间

    if (proposeNum <= 0) {
        LOG_INFO_FMT("No segments to allocate, "
                "proposeNum is {}", proposeNum);
        return KVC_OK;
    }


    int32_t needEvictNum = std::max(proposeNum 
                - (int32_t)_freeList.size(), 0);

    //! 需要驱逐一定数量的Segment
    if (needEvictNum > 0) {
        if (policy == ALLOC_POLICY_DEFAULT) {
            LOG_ERROR_FMT("allocateSegment failed, alloc({}) > free({})", 
                        proposeNum, _freeList.size());
            return KVC_ERR_NOMEM;
        }

        int32_t evictedNum;

        // TODO(yangxianpku: 2025.08.18): 瑕疵，驱逐了部分未恢复
        _lruMap.evict(_segList, _freeList, needEvictNum, evictedNum);
        if (evictedNum < needEvictNum) {
            LOG_ERROR_FMT("evictable segments not enough, alloc({}) > free({})", 
                        proposeNum, _freeList.size());
            return KVC_ERR_NOMEM;
        }
    }

    int32_t allocatedNum = 0;
    for (int32_t i = blockHashIds.size() - 1; i >= 0; i--) {
        // 如果是重复的segmentId，不再新申请
        // TODO: 
        if (segments[i] != nullptr) {
            LOG_INFO_FMT("Segment already allocated for blockHashId: {}",
                        blockHashIds[i].hashId);
            continue;
        }

        uint64_t offset = _freeList.front();

        int32_t idx     = getIdxByOffset(offset);
        assert(offset == _segList[idx].getOffset());
        _freeList.pop_front();

        int64_t segmentId = blockHashIds[i].hashId;
        _segList[idx].setSegmentId(segmentId);
        _segList[idx].setAccessTime(getCurrentTimestampUs());
        _segList[idx].setWriteInRecord(writeInRecordSize);
        _segList[idx].setCacheInfo(cacheInfo);
        _segList[idx].setBlockHashId(blockHashIds[i]);
        _segList[idx].setIsFilled(false);
        _segList[idx].setIsSaved(false);

        // TODO_TLB: cannot use after switching to hugepage, need fix
        // mprotect((void *)((char *)_memStartAddr + _segList[idx].getOffset()), _memSegSize, PROT_READ | PROT_WRITE);
        _lruMap.insert(segmentId, offset, idx);
        segments[i] = &_segList[idx];
        allocatedNum++;
    }
    assert(allocatedNum == proposeNum);

    auto end = getCurrentTimestampUs();  // 记录函数结束时间
    LOG_DEBUG_FMT("cost time : {} us", end - start);
    return KVC_OK;
}



int32_t MemoryPool::allocateSegments(std::vector<BlockHashId> &blockHashIds, 
                                    AllocPolicy policy, 
                                    std::vector<Segment *> &segments,
                                    int32_t writeInRecordSize, 
                                    CacheInfo cacheInfo) {
    auto start = getCurrentTimestampUs();  // 记录函数开始时间

    if (blockHashIds.size() == 0) {
        LOG_INFO_FMT("segmentIds.size() is 0");
        return KVC_OK;
    }

    if (blockHashIds.size() > _blkNum) {
        LOG_ERROR_FMT("allocateSegment failed, alloc({}) > total({})", 
                    blockHashIds.size(), _blkNum);
        return KVC_ERR_INVALID_ARG;
    }

    std::lock_guard<std::shared_mutex> guard(_memPoolLock);

    // 首先，对于重复的hashId，返回已分配的segment
    int32_t proposeNum = 0;
    for (auto blockHashId: blockHashIds) {
        if (_lruMap.contain(blockHashId.hashId)) {
            int32_t idx = _lruMap.getIdxBySegmentId(blockHashId.hashId);

            segments.push_back(&_segList[idx]);
            LOG_INFO_FMT("Segment already allocated for blockHashId: {}", 
                        blockHashId.hashId);
        } else {
            segments.push_back(nullptr);
            proposeNum++;
        }
    }

    // 然后，对于不重复的hashId，新分配segment
    int32_t ret = allocateSegmentByList(blockHashIds, 
                                        policy, 
                                        segments, 
                                        proposeNum,
                                        writeInRecordSize, 
                                        cacheInfo
                                    );
    if (ret != KVC_OK) {
        LOG_ERROR_FMT("allocateSegments failed, ret is {}", ret);
        return ret;
    }

    auto end = getCurrentTimestampUs();  // 记录函数结束时间
    LOG_DEBUG_FMT("cost time : {} us", end - start);
    return KVC_OK;
}



int32_t MemoryPool::sealSegment(int64_t segmentId) {
    std::lock_guard<std::shared_mutex> guard(_memPoolLock);
    auto start = getCurrentTimestampUs();  // 记录函数开始时间

    if (!_lruMap.contain(segmentId)) {
        LOG_DEBUG_FMT("segment not found, segmentId: {}", segmentId);
        return KVC_ERR_INTERNAL;
    }

    int32_t idx = _lruMap.getIdxBySegmentId(segmentId);
    if (_segList[idx].getWriteInRecord() <= 0) {
        LOG_DEBUG_FMT("sealing already sealed segmentId: {}", segmentId);
        return KVC_ERR_INTERNAL;
    }

    // 
    _segList[idx].decWriteInRecord();
    if (_segList[idx].getWriteInRecord() == 0) {
        _segList[idx].setIsFilled(true);
    }

    auto end = getCurrentTimestampUs();  // 记录函数结束时间
    LOG_DEBUG_FMT("sealSegment {} cost time : {} us", 
                segmentId, end - start);
    
    return KVC_OK;
}


int32_t MemoryPool::getSegmentByID(int64_t segmentId, 
                            Segment **pSegmentPtr) {

    std::shared_lock<std::shared_mutex> guard(_memPoolLock);
    if (!_lruMap.contain(segmentId)) {
        LOG_WARN_FMT("segment not found, segmentId: {}", segmentId);
        return KVC_ERR_INTERNAL;
    }

    int32_t idx = _lruMap.getIdxBySegmentId(segmentId);
    if (!_segList[idx].getIsFilled()) {
        LOG_ERROR_FMT("segment not filled, segmentId: {}", segmentId);
        return KVC_ERR_NOT_FILLED;
    }
    *pSegmentPtr = &_segList[idx];

    return KVC_OK;
}



bool MemoryPool::isInMemoryPool(int64_t segmentId)
{
    std::shared_lock<std::shared_mutex> guard(_memPoolLock);
    bool exists = _lruMap.contain(segmentId);
    LOG_INFO_FMT("segmentId: {}, exists: {}", segmentId, exists);
    return exists;
}



std::vector<Segment> MemoryPool::readInOrder(std::vector<int64_t> &segmentIds) {
    std::vector<Segment> segments;

    std::shared_lock<std::shared_mutex> guard(_memPoolLock);
    for (auto &segmentId : segmentIds) {
        if (!_lruMap.contain(segmentId)) {
            LOG_INFO_FMT("segment not found, segmentId: {}", segmentId);
            break;
        }

        int32_t idx = _lruMap.getIdxBySegmentId(segmentId);
        if (!_segList[idx].getIsFilled()) {
            LOG_INFO_FMT("segment not filled, segmentId: {}", segmentId);
            break;
        }

        segments.push_back(_segList[idx]);
    }
    return segments;
}




std::vector<BlockHashId> MemoryPool::getAllCachedSegmentIds() {
    std::shared_lock<std::shared_mutex> guard(_memPoolLock);
    return _lruMap.getAllFilledSegmentIds(_segList);
}


void MemoryPool::releaseSegment(int64_t segmentId) {
    std::lock_guard<std::shared_mutex> guard(_memPoolLock);
    if (!_lruMap.contain(segmentId)) {
        LOG_INFO_FMT("segment not found, segmentId: {}", segmentId);
        return;
    }

    int32_t idx = _lruMap.getIdxBySegmentId(segmentId);
    _segList[idx].setSegmentId(0);
    _segList[idx].setIsFilled(false);
    _segList[idx].setIsSaved(false);
    _segList[idx].setAccessTime(0);
    _lruMap.remove(segmentId);

    _freeList.push_back(_segList[idx].getOffset());
    LOG_INFO_FMT("segmentId: {} released", segmentId);
    return;
}


/**
 *! @brief 清理指定目录下符合特定前缀名的共享内存文件
 * 
 * @param dir_path 文件所在路径
 * @param prefix  指定前缀名的文件
 */
void CleanShmFile(const std::string& dir_path, 
                const std::string& prefix) {
    DIR *dir = opendir(dir_path.c_str());
    if (!dir) {
        LOG_ERROR_FMT("opendir fail! dir_path({})", 
                    dir_path.c_str());
        return;
    }

    int deleted_count = 0;
    struct dirent *entry;
    
    while ((entry = readdir(dir)) != nullptr) {
        if (strcmp(entry->d_name, ".") == 0 || 
            strcmp(entry->d_name, "..") == 0)
            continue;
        
        std::string full_path = dir_path + "/" + entry->d_name;
        
        struct stat path_stat;
        if (stat(full_path.c_str(), &path_stat) != 0) continue;
        if (!S_ISREG(path_stat.st_mode)) continue;
        
        if (strncmp(entry->d_name, prefix.c_str(), prefix.size()) == 0) {
            if (unlink(full_path.c_str()) == 0) {
                LOG_DEBUG_FMT("unlink success! path({})", full_path.c_str());
                deleted_count++;
            } else {
                LOG_ERROR_FMT("unlink failed! path({})", full_path.c_str());
            }
        }
    }
    
    closedir(dir);
    LOG_INFO_FMT("total remove Shm file : {}", deleted_count);
}



/**
 *! @brief 删除共享内存文件
 * 
 * @param shareDir   共享内存文件所在路径
 * @param sharePrefix  共享内存文件前缀
 */
void CleanPortSpecificShmFiles(const std::string& shareDir, 
                        const std::string& sharePrefix) {
    std::string portStr = generateClientPort();
    if (portStr.empty()) {
        LOG_WARN_FMT("Port string is empty, skipping shared file cleanup");
        return;
    }

    DIR* dir = opendir(shareDir.c_str());
    if (!dir) {
        LOG_ERROR_FMT("Failed to open shared memory directory: ", 
                    strerror(errno));
        return;
    }

    std::string targetPrefix = sharePrefix + portStr + "-";

    struct dirent* entry;
    while ((entry = readdir(dir)) != nullptr) {
        if (strcmp(entry->d_name, ".") == 0 || 
            strcmp(entry->d_name, "..") == 0) {
            continue;
        }

        std::string filename = entry->d_name;
        
        if (filename.find(targetPrefix) == 0) {
            std::string fullpath = shareDir + filename;
            if (unlink(fullpath.c_str()) == 0) {
                LOG_DEBUG_FMT("Deleted old shared file: ", 
                            fullpath);
            } else {
                LOG_WARN_FMT("Failed to delete shared file : ",
                            fullpath, " : ", strerror(errno));
            }
        }
    }
    closedir(dir);
}


/**
 *! @brief 打开(创建， 如果存在则删除)共享内存文件返回其文件句柄
 * 
 * @param phyPoolSize   物理内存池大小
 * @param memSegSize    逻辑内存块大小
 * @param pathname      共享内存文件路径
 * @param outFd         共享内存文件句柄
 * 
 * @return int32_t 
 */
int32_t getShareFileFd(uint64_t phyPoolSize, 
                    uint32_t memSegSize, 
                    std::string &pathname, 
                    int32_t &outFd) {
    std::string shareDir    = generateSharedDir();
    std::string sharePrefix = generateShareFilePrefix();

    CleanPortSpecificShmFiles(shareDir, sharePrefix);
    // 打开共享文件
    pathname = generateSharedFilename(memSegSize);

    if (access(pathname.c_str(), F_OK) == 0) {
        // 删除文件
        if (unlink(pathname.c_str()) == 0) {
            LOG_DEBUG_FMT("unlink success!");
        } else {
            LOG_DEBUG_FMT("unlink fail errno({})",strerror(errno));
        }
    }

    int32_t fd = open(pathname.c_str(), O_CREAT | O_RDWR, 0666);
    if (fd < 0) {
        LOG_ERROR_FMT("open share file failed! fd: {}", fd);
        return KVC_ERR_FILEOP_FAIL;
    }

    struct stat statInfo;
    int32_t ret = fstat(fd, &statInfo);
    if (ret != 0) {
        LOG_ERROR_FMT("get file stat failed! file: {}, ret: {}", 
                    pathname.c_str(), ret);
        close(fd);
        return KVC_ERR_FILEOP_FAIL;
    }

    // 如果文件大小不够，则扩展文件大小
    LOG_DEBUG_FMT("phyPoolSize {}", phyPoolSize);
    if (statInfo.st_size < phyPoolSize) {
        ret = ftruncate(fd, phyPoolSize);
        if (ret != 0) {
            LOG_ERROR_FMT("truncate file failed! file: {}, ret: {}", 
                        pathname.c_str(), ret);
            close(fd);
            return KVC_ERR_FILEOP_FAIL;
        }
    }

    outFd = fd;
    return KVC_OK;
}




int32_t createMemoryPool(uint64_t poolSize, 
                        uint32_t memSegSize, 
                        bool isShare, 
                        bool needLock, 
                        MemoryPool *pool) { // pool 
    auto start = getCurrentTimestampUs();  // 记录函数开始时间

	uint64_t phyMemSegSize = ROUND_UP(memSegSize, PAGE_SIZE);
    if (memSegSize == 0) {
        LOG_ERROR_FMT("can't create memory pool, memSegSize is 0!!!");
        return KVC_ERR_FILEOP_FAIL;
    }

    // TODO(yangxianpku: 2025.08.18)： 整除
    int32_t blkNum       = poolSize / memSegSize;
    uint64_t phyPoolSize = ROUND_UP(static_cast<uint64_t>(blkNum) * \
                                    phyMemSegSize, HUGE_PAGE_SIZE);
    
    LOG_DEBUG_FMT("memSegSize:{} phyMemSegSize:{} poolSize:{} phyPoolSize:{}", 
                memSegSize, phyMemSegSize, poolSize, phyPoolSize);
    
    int32_t fd;
    std::string pathname;
    int32_t ret = getShareFileFd(phyPoolSize, memSegSize, pathname, fd);


    if (ret != KVC_OK) {
        LOG_ERROR_FMT("get share file fd failed!");
        return KVC_ERR_FILEOP_FAIL;
    }

    int mapFlags = MAP_SHARED     |   // 映射的内存区域可以被多个进程共享
                    MAP_HUGETLB   |   // 使用大页（Huge Pages）分配内存
                    MAP_POPULATE;     // 在mmap调用时立即分配物理内存
    if (needLock) {
      mapFlags |= MAP_LOCKED;         // 将映射的内存区域 锁定在物理内存中，禁止被换出到磁盘（Swap）
    }
    uint8_t *memStartAddr = (uint8_t *)mmap(NULL,                   // 内核自动选择映射的虚拟地址
                                            phyPoolSize,            // 映射区域大小
                                            PROT_READ | PROT_WRITE, // 映射区域权限
                                            mapFlags,               // 映射标签参数
                                            fd,                     // 共享内存文件描述符
                                            0                       // 文件偏移量
                                        );
    close(fd);
    if (memStartAddr == MAP_FAILED) {
        LOG_ERROR_FMT("mmap failed! file: {}, phyPoolSize: {}", 
                    pathname.c_str(), phyPoolSize);
        
        return KVC_ERR_FILEOP_FAIL;
    }

    uint32_t segmentNum = poolSize / memSegSize;
    std::list<uint64_t> freeList;
    std::vector<Segment> segList;

    for (uint32_t i = 0; i < segmentNum; i++) {
        freeList.push_back(i * phyMemSegSize);
        segList.emplace_back(0, i * phyMemSegSize, 
                    memStartAddr + i * phyMemSegSize, 
                    memSegSize, 
                    false, pathname);
    }

    pool->init(memStartAddr, 
            memSegSize, 
            phyMemSegSize,
            poolSize, 
            phyPoolSize, 
            segmentNum, 
            true, 
            freeList,
            pathname,
            segList
        );

    auto end = getCurrentTimestampUs();  // 记录函数结束时间
    LOG_DEBUG_FMT("cost time : {} us", end - start);
    return KVC_OK;
}




void destroyMemoryPool(MemoryPool* pool) {
    if (pool == nullptr || pool->getMemStartAddr() == nullptr) {
        LOG_INFO_FMT("MemoryPool is null or not initialized");
        return;
    }

    if (munmap(pool->getMemStartAddr(), pool->getPhyPoolSize()) == 0) {
        if (unlink(pool->getShmFilePath().c_str()) != 0) {
            LOG_ERROR_FMT("unlink {} failed!", pool->getShmFilePath().c_str());
        }
    } else {
        LOG_ERROR_FMT("munmap failed!");
    }
    LOG_INFO_FMT("MemoryPool destroyed");
}


std::string MemoryPool::getShmFilePath() {
    return _shmFilePath;
}


uint32_t MemoryPool::getMemSegSize() {
    return _memSegSize;
}


void* MemoryPool::getMemStartAddr() {
    return _memStartAddr;
}


void MemoryPool::reorder(std::vector<int64_t> &segmentIds)
{
    std::lock_guard<std::shared_mutex> guard(_memPoolLock);
    _lruMap.reorder(segmentIds);
    LOG_INFO_FMT("segments reordered");
    return;
}


void MemoryPool::printLruStatus()
{
    _lruMap.print(_segList);
}


void MemoryPool::resetCacheStatis() {
    _statis.dramHit = 0;
    _statis.total   = 0;
    LOG_INFO_FMT("Cache statistics reset");
}


void MemoryPool::updateCacheStatis(int32_t dramHitDelta, 
                                int32_t totalDelta) {
    _statis.dramHit += dramHitDelta;
    _statis.total   += totalDelta;
}


void MemoryPool::printCacheStatis() {
    LOG_DEBUG_FMT("dram hit: {}, total: {}", 
            _statis.dramHit, _statis.total);
}


void MemoryPool::cleanMemPool() {
    std::lock_guard<std::shared_mutex> guard(_memPoolLock);
    _lruMap.clean();
    _segList.clear();
    _freeList.clear();

    uint32_t segmentNum = _poolSize / _memSegSize;
    for (uint32_t i = 0; i < segmentNum; i++) {
        _freeList.push_back(i * _phyMemSegSize);
        _segList.emplace_back(0, i * _phyMemSegSize, 
                        (uint8_t *)_memStartAddr + i * _phyMemSegSize, 
                        _memSegSize, false, _shmFilePath);
    }
    LOG_INFO_FMT("MemoryPool cleaned");
}


uint64_t MemoryPool::getPhyPoolSize() {
    return _phyPoolSize;
}